#include <Winsock2.h>
#include <cstdio>
#include <cstdlib>
#include <cassert>
#include <iostream>
#include <string>
using namespace std;

#pragma  comment(lib,"ws2_32.lib")

int init_win_socket()
{
	WSADATA wsaData;

	if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
	{
		return -1;
	}

	return 0;
}

#define Server_Port 10286
#define MAX_LINE 16384

#define  FD_SETSIZE 1024

struct fd_state
{
	char buffer[MAX_LINE];
	size_t buffer_used;

	int writing;
	size_t n_written;
	size_t write_upto;
};

struct fd_state * alloc_fd_state(void)
{
	struct fd_state *state = (struct fd_state *)malloc(sizeof(struct fd_state));
	if (!state)
		return NULL;
	state->buffer_used = state->n_written = state->writing =
		state->write_upto = 0;
	memset(state->buffer, 0, MAX_LINE);
	return state;
}

void free_fd_state(struct fd_state *state)
{
	free(state);
}

int set_socket_nonblocking(int fd)
{
	unsigned long mode = 1;
	int result = ioctlsocket(fd, FIONBIO, &mode);
	if (result != 0)
	{
		return -1;
		printf("ioctlsocket failed with error: %ld\n", result);
	}
	return 0;
}

int do_read(int fd, struct fd_state *state)
{
	char buf[1024];
	int i;
	int result;
	while (1)
	{
		memset(buf, 0, 1024);
		result = recv(fd, buf, sizeof(buf), 0);
		if (result <= 0)
			break;

		for (i = 0; i < result; ++i)
		{
			if (state->buffer_used < sizeof(state->buffer))
				state->buffer[state->buffer_used++] = buf[i];
		}
	}
	state->writing = 1;
	state->write_upto = state->buffer_used;
	printf("Receive data: %s  size: %d\n", state->buffer + state->n_written, state->write_upto - state->n_written);

	if (result == 0)
	{
		return 1;
	}
	else if (result < 0)
	{
#ifdef WIN32
		if (result == -1 && WSAGetLastError() == WSAEWOULDBLOCK)
			return 0;
#else
		if (errno == EAGAIN)
			return 0;
#endif
		return -1;
	}

	return 0;
}

int do_write(int fd, struct fd_state *state)
{
	while (state->n_written < state->write_upto)
	{
		int result = send(fd, state->buffer + state->n_written,
			state->write_upto - state->n_written, 0);
		if (result < 0)
		{
#ifdef WIN32
			if (result == -1 && WSAGetLastError() == WSAEWOULDBLOCK)
				return 0;
#else
			if (errno == EAGAIN)
				return 0;
#endif
			return -1;
		}
		assert(result != 0);
		printf("Send data: %s \n", state->buffer + state->n_written);
		state->n_written += result;
	}

	if (state->n_written == state->buffer_used)
		state->n_written = state->write_upto = state->buffer_used = 0;

	state->writing = 0;

	return 0;
}

void run()
{
	int listener;
	struct fd_state *state[FD_SETSIZE];
	struct sockaddr_in sin;
	int i, maxfd;
	fd_set readset, writeset, exset;

	sin.sin_family = AF_INET;
	sin.sin_addr.s_addr = 0;
	sin.sin_port = htons(Server_Port);

	for (i = 0; i < FD_SETSIZE; ++i)
		state[i] = NULL;

	listener = socket(AF_INET, SOCK_STREAM, 0);
	set_socket_nonblocking(listener);

	int one = 1;
	setsockopt(listener, SOL_SOCKET, SO_REUSEADDR, (const char *)&one, sizeof(one));

	if (bind(listener, (struct sockaddr*)&sin, sizeof(sin)) < 0)
	{
		perror("bind");
		return;
	}

	if (listen(listener, 16)<0)
	{
		perror("listen");
		return;
	}

	printf("Server is listening ... \n");

	FD_ZERO(&readset);
	FD_ZERO(&writeset);
	FD_ZERO(&exset);

	while (1)
	{
		maxfd = listener;

		FD_ZERO(&readset);
		FD_ZERO(&writeset);
		FD_ZERO(&exset);

		FD_SET(listener, &readset);

		for (i = 0; i < FD_SETSIZE; ++i)
		{
			if (state[i])
			{
				if (i > maxfd)
					maxfd = i;
				FD_SET(i, &readset);
				if (state[i]->writing)
				{
					FD_SET(i, &writeset);
				}
			}
		}

		if (select(maxfd + 1, &readset, &writeset, &exset, NULL) < 0)
		{
			perror("select");
			return;
		}

		//check if listener can accept
		if (FD_ISSET(listener, &readset))
		{
			struct sockaddr_in ss;
			int slen = sizeof(ss);
			int fd = accept(listener, (struct sockaddr*)&ss, &slen);
			if (fd < 0)
			{
				perror("accept");
			}
			else if (fd > FD_SETSIZE)
			{
				closesocket(fd);
			}
			else
			{
				printf("Accept socket %d, address %s \n", fd, inet_ntoa(ss.sin_addr));
				set_socket_nonblocking(fd);
				state[fd] = alloc_fd_state();
				assert(state[fd]);
			}
		}

		//process read and write socket
		for (i = 0; i < maxfd + 1; ++i)
		{
			int r = 0;
			if (i == listener)
				continue;

			if (FD_ISSET(i, &readset))
			{
				r = do_read(i, state[i]);
			}
			if (r == 0 && FD_ISSET(i, &writeset))
			{
				r = do_write(i, state[i]);
			}
			if (r)
			{
				free_fd_state(state[i]);
				state[i] = NULL;
				closesocket(i);
			}
		}
	}
}

int main(int c, char **v)
{
#ifdef WIN32
	init_win_socket();
#endif

	run();
	return 0;
}